# include "scene.h"

#ifdef SOLVE_ENERGY

# include "discr.h"
# include <stdio.h>
# include <math.h>
# include <stdlib.h>

#ifdef _OMP_TEMP_ADVANCE_
# include <omp.h>
#endif

# ifdef THMCOND_TDEP
extern double thm_cond(double T);
# endif

extern void temp_bc(struct grid *g, double ***T0_imp, short tflag, unsigned short iRK);
// extern void temp_bc(const struct grid *g, double ****T, short tflag, unsigned short iRK/*, double *pE_accum*/);

static double getE(struct grid *g, double ***T0, unsigned short i, unsigned short j, unsigned short k, unsigned short isec);
static unsigned int factorial(unsigned short n);

/* Solving the Energy Equation using 4-stage RK-3 scheme. */
extern double temp_advance(struct grid *g, const struct arr_struct *depVar)
  {
   short i, j, k, isec, impsec;
   double ***_T[Nsec][5], *r, *th, *z, E_absr;

   for(isec = 0; isec < Nsec; isec++)
     {/* Assigning space for different RK-stages. */
      _T[isec][0] = g->T[isec];
      _T[isec][1] = depVar[isec].q[7]; _T[isec][3] = depVar[isec].q[7];
      _T[isec][2] = depVar[isec].q[8]; _T[isec][4] = depVar[isec].q[8];
     }

/* The section in which impingement on the plate is taking place. */
   impsec = 1;

/* If temperature in the impingement-section reaches the expected temperature, turn-off the volume heat source. */
   if(g->vhsrc_flag == 'y')
   if( _T[impsec][0][NrG+2*g->nr[impsec]/3][NthG+g->nth/2][NzG+g->nz[impsec]/2]-T_INLET > EXPECTED_T_RISE) g->vhsrc_flag = 'n';


/* ------------------- Calculating T(1) ------------------- */
   for(isec = 0; isec < Nsec; isec++)
       {
#       ifdef _OMP_TEMP_ADVANCE_
        #pragma omp parallel for private(i,j,k) schedule(dynamic,1)
#       endif
        for(i = NrG; i < g->nr[isec]+NrG; i++) for(j = NthG; j < g->nth+NthG; j++) for(k = NzG; k < g->nz[isec]+NzG; k++)
           _T[isec][1][i][j][k] = _T[isec][0][i][j][k] + 0.5*g->dt*getE(g,_T[isec][0],i,j,k,isec);

        g->T[isec] = _T[isec][1];
       }
   temp_bc(g, _T[impsec][0], 0, 1);
/* -------------------------------------------------------- */

/* ------------------- Calculating T(2) ------------------- */
   for(isec = 0; isec < Nsec; isec++)
       {
#       ifdef _OMP_TEMP_ADVANCE_
        #pragma omp parallel for private(i,j,k) schedule(dynamic,1)
#       endif
        for(i = NrG; i < g->nr[isec]+NrG; i++) for(j = NthG; j < g->nth+NthG; j++) for(k = NzG; k < g->nz[isec]+NzG; k++)
           _T[isec][2][i][j][k] = _T[isec][1][i][j][k] + 0.5*g->dt*getE(g,_T[isec][0],i,j,k,isec);

        g->T[isec] = _T[isec][2];
       }
   temp_bc(g, _T[impsec][0], 0, 2);
/* -------------------------------------------------------- */

/* ------------------- Calculating T(3) ------------------- */
   for(isec = 0; isec < Nsec; isec++)
       {
#       ifdef _OMP_TEMP_ADVANCE_
        #pragma omp parallel for private(i,j,k) schedule(dynamic,1)
#       endif
        for(i = NrG; i < g->nr[isec]+NrG; i++) for(j = NthG; j < g->nth+NthG; j++) for(k = NzG; k < g->nz[isec]+NzG; k++)
          _T[isec][3][i][j][k] = 2*_T[isec][0][i][j][k]/3 + _T[isec][2][i][j][k]/3 + g->dt*getE(g,_T[isec][0],i,j,k,isec)/6;

        g->T[isec] = _T[isec][3];
       }
   temp_bc(g, _T[impsec][0], 0, 3);
/* -------------------------------------------------------- */

/* ------- Getting "T" at next time-step. It will be given by "_T[4]" --------- */
   for(isec = 0; isec < Nsec; isec++)
       {
#       ifdef _OMP_TEMP_ADVANCE_
        #pragma omp parallel for private(i,j,k) schedule(dynamic,1)
#       endif
        for(i = NrG; i < g->nr[isec]+NrG; i++) for(j = NthG; j < g->nth+NthG; j++) for(k = NzG; k < g->nz[isec]+NzG; k++)
          _T[isec][4][i][j][k] = _T[isec][3][i][j][k] + 0.5*g->dt*getE(g,_T[isec][0],i,j,k,isec);

        g->T[isec] = _T[isec][4];
       }
   temp_bc(g, _T[impsec][0], 1, 0);
/* ------------------------------------------------------------------------------ */

/* Determining total energy absorbed by the system. */
   for(E_absr = 0.0, isec = 0; isec < Nsec; isec++)
      {
       r = g->r[isec]; z = g->z[isec]; th = g->th;

#      ifdef _OMP_TEMP_ADVANCE_
       #pragma omp parallel private(i,j,k)
#      endif
          {
#          ifdef _OMP_TEMP_ADVANCE_
           #pragma omp for schedule(dynamic,3) reduction(+:E_absr)
#          endif
           for(i = NrG; i < g->nr[isec]+NrG; i++) for(j = NthG; j < g->nth+NthG; j++) for(k = NzG+g->nz[isec]-1; k >= NzG; k--)
           E_absr += (r[i+1]+r[i])*(r[i+1]-r[i])*(th[j+1]-th[j])*(z[k+1]-z[k])*(_T[isec][4][i][j][k] - _T[isec][0][i][j][k]);

#          ifdef _OMP_TEMP_ADVANCE_
           #pragma omp for schedule(dynamic,3)
#          endif
           for(i = 0; i < g->nr[isec]+2*NrG; i++) for(j = 0; j < g->nth+2*NthG; j++) for(k = 0; k < g->nz[isec]+2*NzG; k++)
           _T[isec][0][i][j][k] = _T[isec][4][i][j][k]; /* Update the temperature with new value. */
          }

       g->T[isec] = _T[isec][0];
      }

   E_absr *= 0.5*g->density0*Cp_WATER;

   return E_absr;
  }


static double getE(struct grid *g, double ***T0, unsigned short i, unsigned short j, unsigned short k, unsigned short isec)
  {
   double ***u, ***v, ***w, ***T, *r, *th, *z, thmcond0, density0;
   struct cmtrx *cm;
   double q1, q2, q3, conv_term, conduct_term;
   char vhsrc_flag;
   unsigned short indx;
   unsigned int factorialn;

   u = g->um[isec];  v = g->vm[isec];  w = g->wm[isec]; T = g->T[isec];  r = g->r[isec];  th = g->th;  z = g->z[isec];
   thmcond0 = g->thmcond0; density0 = g->density0;

   vhsrc_flag = g->vhsrc_flag;

#  ifndef DISCR_AUTO
   return EE;
#  endif

   factorialn = factorial(nDISCR);
   conv_term = 0.0;
   conduct_term = 0.0;

/* Evaluating u*dT/dr at the cell-center (CD+upwind). */
// Calculating u at the cell-center.
   cm = &g->icmx_rsc[isec];
   for(q1 = 0.0, indx = 0; indx <= cm->iend[i-cm->respadi]-cm->istart[i-cm->respadi]; indx++)
   q1 += cm->mtrx[i-cm->respadi][indx]*u[cm->srcpadi-1+cm->istart[i-cm->respadi]+indx][j][k];

// Calculating dT/dr at the cell-center.
   cm = &g->dcmx_rcc1[isec];
   for(q2 = 0.0, indx = 0; indx <= cm->iend[i-cm->respadi]-cm->istart[i-cm->respadi]; indx++)
     q2 += cm->mtrx[i-cm->respadi][indx]*T[cm->srcpadi+cm->istart[i-cm->respadi]+indx][j][k];

/* Calculating (1/r)*(dT/dr) at the cell-center. */
   conduct_term += 2*q2/(r[i]+r[i+1]);

// Calculating dmT/drm at the cell-center. Here, m = nDISCR; It is 2nd order accurate discretization.
   cm = &g->dcmx_rccm[isec];
   for(q3 = 0.0, indx = 0; indx <= cm->iend[i-cm->respadi]-cm->istart[i-cm->respadi]; indx++)
     q3 += cm->mtrx[i-cm->respadi][indx]*T[cm->srcpadi+cm->istart[i-cm->respadi]+indx][j][k];

   conv_term += q1*( q2\
                     + (2*EEQ_DISSIP_FACT*(q1>0 ? -r[i+1]+r[i-1]:r[i+2]-r[i])/(r[i+2]+r[i+1]-r[i]-r[i-1]))\
                       *(pow((r[i+2]+r[i+1]-r[i]-r[i-1])/4, nDISCR-1)/factorialn)\
                       *q3 );


/* Evaluating (v/r)*(dt/dth) using hybrid scheme (CD+upwind). */
// Calculating v at the cell-center.
   cm = &g->icmx_thsc;
   for(q1 = 0.0, indx = 0; indx <= cm->iend[j-cm->respadi]-cm->istart[j-cm->respadi]; indx++)
   q1 += cm->mtrx[j-cm->respadi][indx]*v[i][cm->srcpadi-1+cm->istart[j-cm->respadi]+indx][k];

// Calculating dT/dth at the cell-center.
   cm = &g->dcmx_thcc1;
   for(q2 = 0.0, indx = 0; indx <= cm->iend[j-cm->respadi]-cm->istart[j-cm->respadi]; indx++)
     q2 += cm->mtrx[j-cm->respadi][indx]*T[i][cm->srcpadi+cm->istart[j-cm->respadi]+indx][k];

// Calculating dmT/dthm at the cell-center. Here, m = nDISCR; It is 2nd-order accurate discretization.
   cm = &g->dcmx_thccm;
   for(q3 = 0.0, indx = 0; indx <= cm->iend[j-cm->respadi]-cm->istart[j-cm->respadi]; indx++)
     q3 += cm->mtrx[j-cm->respadi][indx]*T[i][cm->srcpadi+cm->istart[j-cm->respadi]+indx][k];

   conv_term += (2*q1/(r[i+1]+r[i]))*( q2\
                        + (2*EEQ_DISSIP_FACT*(q1>0 ? -th[j+1]+th[j-1]:th[j+2]-th[j])/(th[j+2]+th[j+1]-th[j]-th[j-1]))\
                          *(pow((th[j+2]+th[j+1]-th[j]-th[j-1])/4, nDISCR-1)/factorialn)\
                          *q3 );


/* Evaluating w*(dT/dz) using hybrid scheme (CD+upwind). */
// Calculating w at the cell-center.
   cm = &g->icmx_zsc[isec][0];
   if(isec == 1) if( (i >= NrG+g->nr[0]) && (i < NrG+g->nr[1]-g->nr[2])) cm = &g->icmx_zsc[isec][1];
   for(q1 = 0.0, indx = 0; indx <= cm->iend[k-cm->respadi]-cm->istart[k-cm->respadi]; indx++)
   q1 += cm->mtrx[k-cm->respadi][indx]*w[i][j][cm->srcpadi-1+cm->istart[k-cm->respadi]+indx];

// Calculating dT/dz at the cell-center.
   cm = &g->dcmx_zcc1[isec][0];
   if(isec == 1) if( (i >= NrG+g->nr[0]) && (i < NrG+g->nr[1]-g->nr[2])) cm = &g->dcmx_zcc1[isec][1];
   for(q2 = 0.0, indx = 0; indx <= cm->iend[k-cm->respadi]-cm->istart[k-cm->respadi]; indx++)
     q2 += cm->mtrx[k-cm->respadi][indx]*T[i][j][cm->srcpadi+cm->istart[k-cm->respadi]+indx];

// Calculating dmT/dzm at the cell-center. Here, m = nDISCR; It is 2nd-order accurate discretization.
   cm = &g->dcmx_zccm[isec][0];
   if(isec == 1) if( (i >= NrG+g->nr[0]) && (i < NrG+g->nr[1]-g->nr[2])) cm = &g->dcmx_zccm[isec][1];
   for(q3 = 0.0, indx = 0; indx <= cm->iend[k-cm->respadi]-cm->istart[k-cm->respadi]; indx++)
     q3 += cm->mtrx[k-cm->respadi][indx]*T[i][j][cm->srcpadi+cm->istart[k-cm->respadi]+indx];

   conv_term += q1*( q2\
                      + (2*EEQ_DISSIP_FACT*(q1>0 ? -z[k+1]+z[k-1]:z[k+2]-z[k])/(z[k+2]+z[k+1]-z[k]-z[k-1]))\
                        *(pow((z[k+2]+z[k+1]-z[k]-z[k-1])/4, nDISCR-1)/factorialn)\
                        *q3 );


/* Evaluating d2T/dr2 at the cell-center. */
   cm = &g->dcmx_rcc2[isec];
   for(q1 = 0.0, indx = 0; indx <= cm->iend[i-cm->respadi]-cm->istart[i-cm->respadi]; indx++)
     q1 += cm->mtrx[i-cm->respadi][indx]*T[cm->srcpadi+cm->istart[i-cm->respadi]+indx][j][k];

   conduct_term += q1;

/* Evaluating (1/r^2)*(d2T/dth2) at the cell-center. */
   cm = &g->dcmx_thcc2;
   for(q1 = 0.0, indx = 0; indx <= cm->iend[j-cm->respadi]-cm->istart[j-cm->respadi]; indx++)
     q1 += cm->mtrx[j-cm->respadi][indx]*T[i][cm->srcpadi+cm->istart[j-cm->respadi]+indx][k];

   conduct_term += 4*(q1/(r[i]+r[i+1]))/(r[i]+r[i+1]);

/* Evaluating d2T/dz2 at the cell-center. */
   cm = &g->dcmx_zcc2[isec][0];
   if(isec == 1) if( (i >= NrG+g->nr[0]) && (i < NrG+g->nr[1]-g->nr[2])) cm = &g->dcmx_zcc2[isec][1];
   for(q1 = 0.0, indx = 0; indx <= cm->iend[k-cm->respadi]-cm->istart[k-cm->respadi]; indx++)
     q1 += cm->mtrx[k-cm->respadi][indx]*T[i][j][cm->srcpadi+cm->istart[k-cm->respadi]+indx];

   conduct_term += q1;

   return (vhsrc_flag == 'y' ? EXPECTED_T_RISE/(double)EXPECTED_TIME : 0) + (conduct_term*thmcond0/density0)/Cp_WATER - conv_term;
  }


static unsigned int factorial(unsigned short n)
  {
   unsigned int nfactorial;
   unsigned short i;

   nfactorial = 1;
   for(i = 1; i <= n; i++) nfactorial *= i;

   return nfactorial;
  }

#endif
